package org.example.datastream.test.custom;

import org.apache.commons.beanutils.BeanUtils;
import org.apache.flink.api.common.serialization.DeserializationSchema;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.typeutils.ResultTypeQueryable;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;

import javax.persistence.Column;
import javax.persistence.Entity;
import java.lang.reflect.Field;
import java.sql.*;

/**
 * 自定义通用jdbc source
 * @param <T> source 泛型
 */
public class MyCustomJdbcSource<T> extends RichSourceFunction<T> implements ResultTypeQueryable<T> {

    private PreparedStatement ps;
    private Connection conn;
    private ResultSet rs;

    private String driver;
    private String url;
    private String username;
    private String password;
    private String sql;
    private Class<T> sourceClass;
    private DeserializationSchema<T> schema;

    public MyCustomJdbcSource(String driver, String url, String username, String password,
                              String sql, Class<T> sourceClass, DeserializationSchema<T> schema) {
        this.driver = driver;
        this.url = url;
        this.username = username;
        this.password = password;
        this.sql = sql;
        this.sourceClass = sourceClass;
        this.schema = schema;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        super.open(parameters);

        Class.forName(driver);

        conn = DriverManager.getConnection(url, username, password);
        ps = conn.prepareStatement(sql);
    }

    @Override
    public void run(SourceContext<T> sourceContext) throws Exception {
        rs = ps.executeQuery();
        if (sourceClass.isAnnotationPresent(Entity.class)) {
            while(rs.next()) {
                ResultSetMetaData rsmd = rs.getMetaData();
                T bean = (T) sourceClass.newInstance();
                Field[] fields = sourceClass.getDeclaredFields();
                for (int _ite = 0; _ite < rsmd.getColumnCount(); _ite++) {
                    String columnName = rsmd.getColumnName(_ite + 1);
                    Object columnValue = rs.getObject(_ite + 1);

                    for (Field field : fields) {
                        if (field.isAnnotationPresent(Column.class)) {
                            Column column = field.getAnnotation(Column.class);
                            if (column.name().equalsIgnoreCase(columnName) && columnValue != null) {
                                BeanUtils.setProperty(bean, field.getName(), columnValue);
                                break;
                            }
                        }
                    }
                }
                sourceContext.collect(bean);
            }
        }
    }

    @Override
    public void cancel() {

    }

    @Override
    public void close() throws Exception {
        super.close();

        if (null != conn) {
            conn.close();
        }
        if (null != ps) {
            ps.close();
        }
    }

    @Override
    public TypeInformation<T> getProducedType() {
        return schema.getProducedType();
    }
}
